#!/usr/bin/env python
import numpy as np
from astropy.io import fits
from glob import glob
import sys
import os

for root,dirs,files in os.walk(sys.argv[1]):
    for name in files:
        #print(root.split('/'))
        flag = 0
        if name[-4:] == 'fits':
            if len(root.split('/')) == 1:
                flag = 1
            elif root.split('/')[1] != 'image':
                flag = 1
        if flag == 1: 
            p = os.path.join(root,name)
            print("cropping file "+p)
            hdulist = fits.open(p)
            for i in range(len(hdulist)):
                data = hdulist[i].data
                print("before crop",data.shape)
                header = hdulist[i].header
                if sys.argv[2] == '0083':
                    if data.shape[1] == 4616:
                        print("cropped already")
                        continue
                    data_crop = data.copy()[:4616,:]
                if sys.argv[2] == '0095':
                    if data.shape[1] == 4616:
                        print("cropped already")
                        continue
                    data_crop = data.copy()[4616:,:]
                print("after crop",data_crop.shape)
                if i==0:
                    fits.writeto(p,data=data_crop,header=header,overwrite=True)
                else:
                    fits.append(p,data=data_crop,header=header)
